import numpy as np  # numpy：数据处理常用库（需要下载）
import argparse     # argparse：解析命令行调用脚本时的命令行输入（Python自带）
import os           # os：文件操作（Python自带）
import shutil       # shutil：文件操作（Python自带）

# gen random data in any shape
def gen_random_data_by_shape_and_type(left, right, shape, target_type):
    target = np.random.uniform(left, right, shape).astype(target_type)
    if shape == (1):
        return target[0]
    return target



# 目前只考虑了全行优先的输入
def gen_data_cal_expect_write_bin(wholeM, wholeN, wholeK, batchCount, check):
    
    print("Get input data: wholeM=%d, wholeN=%d, wholeK=%d, batchCount=%d"%(wholeM, wholeN, wholeK, batchCount))
    print("="*10)

    if check:
        print("wholeM: "+str(wholeM)
                +"\twholeN: "+str(wholeN)
                +"\twholeK: "+str(wholeK)
                +"\tbatchCount: "+str(batchCount))

    if os.path.exists("data/"):
        shutil.rmtree("data/")
    os.mkdir("data/")

    alpha = gen_random_data_by_shape_and_type(-1, 1, (1), np.float16)
    alpha.tofile("data/alpha.bin", sep= " " if check else "")
    if check:
        print("alpha: "+str(alpha))

    beta  = gen_random_data_by_shape_and_type(-1, 1, (1), np.float16)
    beta.tofile("data/beta.bin", sep=" " if check else "")
    if check:
        print("beta: "+str(beta))

    with open(f"data/Q.bin", "a" if check else "ab") as Q_file, \
         open(f"data/KT.bin", "a" if check else "ab") as KT_file, \
         open(f"data/C1.bin", "a" if check else "ab") as C1_file, \
         open(f"data/V.bin", "a" if check else "ab") as V_file, \
         open(f"data/C2.bin", "a" if check else "ab") as C2_file, \
         open(f"data/Parameters.bin", "a" if check else "ab") as Parameters_file, \
         open(f"data/C3.bin", "a" if check else "ab") as C3_file, \
         open(f"data/validM.bin", "a" if check else "ab") as validM_file, \
         open(f"data/expect_result1.bin", "a" if check else "ab") as expect_result1_file, \
         open(f"data/expect_result2.bin", "a" if check else "ab") as expect_result2_file, \
         open(f"data/expect_result3.bin", "a" if check else "ab") as expect_result3_file:

        for batchNum in range(batchCount):

            print("BatchNum %d generating..."%(batchNum))

            # np.array([wholeM]).astype(np.int64)[0] # 
            validM = gen_random_data_by_shape_and_type(wholeM/2, wholeM, (1), np.int64)

            if check: 
                print("validM: "+str(validM))
            
            shape_Q_valid = (validM, wholeK)
            shape_KT_valid = (wholeK, validM)
            shape_C1_valid = (validM, validM)  # 同expect_result1_valid shape

            shape_V_valid = (validM, wholeK)
            shape_C2_valid = (validM, wholeK)  # 同expect_result2_valid shape

            shape_Parameters_valid = (wholeK, wholeN)
            shape_C3_valid = (validM, wholeN)   # 同expect_result3_valid shape

            l, r = -0.125, 0.125
            
            Q_valid = gen_random_data_by_shape_and_type(l, r, shape_Q_valid, np.float16)
            KT_valid = gen_random_data_by_shape_and_type(l, r, shape_KT_valid, np.float16)
            C1_valid = gen_random_data_by_shape_and_type(l, r, shape_C1_valid, np.float16 )

            V_valid = gen_random_data_by_shape_and_type(l, r, shape_V_valid, np.float16)
            C2_valid = gen_random_data_by_shape_and_type(l, r, shape_C2_valid, np.float16)

            Parameters_valid = gen_random_data_by_shape_and_type(l, r, shape_Parameters_valid, np.float16)
            C3_valid = gen_random_data_by_shape_and_type(l, r, shape_C3_valid, np.float16)

            print("BatchNum %d validABC generation finished."%(batchNum))

            temp = np.dot(Q_valid.astype(np.float32), KT_valid.astype(np.float32)).astype(np.float16)
            expect_result1_valid = alpha * temp + beta * C1_valid
            temp = np.dot(expect_result1_valid.astype(np.float32), V_valid.astype(np.float32)).astype(np.float16)
            expect_result2_valid = alpha * temp + beta * C2_valid
            temp = np.dot(expect_result2_valid.astype(np.float32), Parameters_valid.astype(np.float32)).astype(np.float16)
            expect_result3_valid = alpha * temp + beta * C3_valid

            print("BatchNum %d validExpectResult generation finished." % (batchNum))

            if check:
                print("expect_result1_valid: ")
                print(expect_result1_valid)
                print("expect_result2_valid: ")
                print(expect_result2_valid)
                print("expect_result3_valid: ")
                print(expect_result3_valid)
            
            Q_zero_padding = np.zeros((wholeM, wholeK), dtype=np.float16)
            KT_zero_padding = np.zeros((wholeK, wholeM), dtype=np.float16)
            C1_zero_padding = np.zeros((wholeM, wholeM), dtype=np.float16)
            expect_result1_zero_padding = np.zeros((wholeM, wholeM), dtype=np.float16)

            V_zero_padding = np.zeros((wholeM, wholeK), dtype=np.float16)
            C2_zero_padding = np.zeros((wholeM, wholeK), dtype=np.float16)
            expect_result2_zero_padding = np.zeros((wholeM, wholeK), dtype=np.float16)

            Parameters_zero_padding = np.zeros((wholeK, wholeN), dtype=np.float16)
            C3_zero_padding = np.zeros((wholeM, wholeN), dtype=np.float16)
            expect_result3_zero_padding = np.zeros((wholeM, wholeN), dtype=np.float16)


            Q_zero_padding[:validM, :wholeK] = Q_valid
            KT_zero_padding[:wholeK, :validM] = KT_valid
            C1_zero_padding[:validM, :validM] = C1_valid
            expect_result1_zero_padding[:validM, :validM] = expect_result1_valid

            V_zero_padding[:validM, :wholeK] = V_valid
            C2_zero_padding[:validM, :wholeK] = C2_valid
            expect_result2_zero_padding[:validM, :wholeK] = expect_result2_valid

            Parameters_zero_padding[:wholeK, :wholeN] = Parameters_valid
            C3_zero_padding[:validM, :wholeN] = C3_valid
            expect_result3_zero_padding[:validM, :wholeN] = expect_result3_valid

            print("BatchNum %d padding data generation finished." % (batchNum))

            if check:
                print("Q_zero_padding: ")
                print(Q_zero_padding)
                print("KT_zero_padding: ")
                print(KT_zero_padding)
                print("expect_result1_zero_padding: ")
                print(expect_result1_zero_padding)
                print("expect_result2_zero_padding: ")
                print(expect_result2_zero_padding)
                print("expect_result3_zero_padding: ")
                print(expect_result3_zero_padding)

            # tofile
            Q_zero_padding.tofile(Q_file, sep=" " if check else "")
            KT_zero_padding.tofile(KT_file, sep=" " if check else "")
            C1_zero_padding.tofile(C1_file, sep=" " if check else "")
            expect_result1_zero_padding.tofile(expect_result1_file, sep=" " if check else "")

            V_zero_padding.tofile(V_file, sep=" " if check else "")
            C2_zero_padding.tofile(C2_file, sep=" " if check else "")
            expect_result2_zero_padding.tofile(expect_result2_file, sep=" " if check else "")

            Parameters_zero_padding.tofile(Parameters_file, sep=" " if check else "")
            C3_zero_padding.tofile(C3_file, sep=" " if check else "")
            expect_result3_zero_padding.tofile(expect_result3_file, sep=" " if check else "")

            validM.tofile(validM_file, sep=" " if check else "")
            
            print("BatchNum %d tofile finished."%(batchNum))
            print("="*10)
            pass  # for batch

        pass  # with open

    pass  # gen_data_cal_expect_write_bin


if __name__ == "__main__":

    parser = argparse.ArgumentParser()
    parser.add_argument("wholeM", action="store", type=int)
    parser.add_argument("wholeN", action="store", type=int)
    parser.add_argument("wholeK", action="store", type=int)
    parser.add_argument("batchCount", action="store", type=int)

    args = parser.parse_args()

    wholeM = args.wholeM
    wholeN = args.wholeN
    wholeK = args.wholeK
    batchCount = args.batchCount

    gen_data_cal_expect_write_bin(wholeM, wholeN, wholeK, batchCount, False)